#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
快速漏洞扫描演示脚本
演示漏洞扫描器的核心功能，但使用更小的扫描范围以提供更快的执行速度
"""

import os
import time
import json
import logging
import requests
import argparse
from datetime import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('quick_scan_demo')

# API基础URL
BASE_URL = "http://localhost:5000"

# 示例目标列表 - 更小、更快的目标
QUICK_TARGETS = [
    {
        "name": "OWASP WebGoat",
        "url": "http://owasp-webgoat.herokuapp.com/WebGoat",
        "description": "OWASP WebGoat项目，专注于Web安全培训",
        "config": {"depth": 1, "scan_types": ["XSS", "CSRF"]}
    },
    {
        "name": "Web Application Security Assessment Framework",
        "url": "http://demo.websecurify.com",
        "description": "Web安全评估框架测试网站",
        "config": {"depth": 1, "scan_types": ["SQL注入", "XSS"]}
    },
    {
        "name": "Hack The Box Challenges",
        "url": "https://www.hackthebox.eu/home/challenges",
        "description": "网络安全挑战平台",
        "config": {"depth": 1, "scan_types": ["XSS", "目录遍历"]}
    }
]

def print_header(title):
    """打印美观的标题"""
    border = "=" * 70
    print(f"\n{border}")
    print(f"{title.center(66)}")
    print(f"{border}\n")

def print_separator():
    """打印分隔线"""
    print("-" * 70)

def check_api_health():
    """检查API服务是否正常运行"""
    try:
        response = requests.get(f"{BASE_URL}/api/health", timeout=5)
        return response.status_code == 200
    except (requests.RequestException, Exception):
        logger.error("API健康检查失败")
        return False

def submit_scan_task(target_url, scan_config):
    """提交扫描任务到API"""
    try:
        payload = {
            "target": target_url,  # 修改为正确的参数名
            "scan_depth": scan_config["depth"],
            "scan_types": scan_config["scan_types"]
        }
        response = requests.post(
            f"{BASE_URL}/api/scan",
            json=payload,
            timeout=10
        )
        
        if response.status_code in [200, 202]:  # 202表示Accepted，也是成功的状态
            data = response.json()
            return data.get("task_id") or data.get("scan_id")
        else:
            logger.error(f"提交扫描任务失败: {response.status_code} - {response.text}")
            return None
    except (requests.RequestException, Exception) as e:
        logger.error(f"提交扫描任务时发生错误: {str(e)}")
        return None

def get_scan_status(task_id):
    """获取扫描任务的状态"""
    try:
        response = requests.get(
            f"{BASE_URL}/api/scan/{task_id}/status",
            timeout=5
        )
        
        if response.status_code == 200:
            return response.json()
        else:
            logger.error(f"获取扫描状态失败: {response.status_code}")
            return {"status": "error", "progress": 0}
    except (requests.RequestException, Exception) as e:
        logger.error(f"获取扫描状态时发生错误: {str(e)}")
        return {"status": "error", "progress": 0}

def get_scan_report(task_id):
    """获取扫描报告"""
    try:
        response = requests.get(
            f"{BASE_URL}/api/scan/{task_id}/report",
            timeout=10
        )
        
        if response.status_code == 200:
            return response.json()
        else:
            logger.error(f"获取扫描报告失败: {response.status_code}")
            return {"status": "error"}
    except (requests.RequestException, Exception) as e:
        logger.error(f"获取扫描报告时发生错误: {str(e)}")
        return {"status": "error"}

def save_scan_result(result, target_name, task_id):
    """保存扫描结果到JSON文件"""
    filename = f"quick_scan_{target_name.replace(' ', '_')}_{task_id}.json"
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        return filename
    except Exception as e:
        logger.error(f"保存扫描结果失败: {str(e)}")
        return None

def perform_quick_scan(target):
    """执行快速扫描"""
    print_header(f"快速扫描: {target['name']}")
    print(f"目标URL: {target['url']}")
    print(f"描述: {target['description']}")
    print(f"扫描配置: 深度={target['config']['depth']}, 类型={', '.join(target['config']['scan_types'])}")
    print_separator()
    
    start_time = time.time()
    logger.info(f"开始扫描目标: {target['name']}, URL: {target['url']}")
    
    # 提交扫描任务
    print("正在提交扫描任务...")
    task_id = submit_scan_task(target['url'], target['config'])
    
    if not task_id:
        print("❌ 扫描任务提交失败")
        return False
    
    print(f"✅ 扫描任务创建成功，任务ID: {task_id}")
    
    # 等待扫描完成 - 设置1分钟超时以避免长时间等待
    max_wait_time = 60  # 秒
    wait_interval = 3  # 秒
    elapsed_time = 0
    
    while elapsed_time < max_wait_time:
        status_data = get_scan_status(task_id)
        status = status_data.get("status", "unknown")
        progress = status_data.get("progress", 0)
        
        print(f"🔄 扫描状态: {status}, 进度: {progress}%")
        
        if status in ["completed", "error", "failed"]:
            break
            
        time.sleep(wait_interval)
        elapsed_time += wait_interval
    
    if elapsed_time >= max_wait_time:
        print("⚠️  扫描超时，尝试获取当前结果")
    
    # 获取扫描报告
    print("正在获取扫描报告...")
    scan_time = time.time() - start_time
    report = get_scan_report(task_id)
    
    # 保存结果
    if report and "status" in report and report["status"] != "error":
        filename = save_scan_result(report, target['name'], task_id)
        print(f"✅ 扫描完成，耗时: {scan_time:.2f}秒")
        print(f"📄 扫描结果已保存到: {filename}")
        
        # 显示简要结果
        display_quick_summary(report)
        return True
    else:
        print("❌ 无法获取扫描报告")
        return False

def display_quick_summary(report):
    """显示快速扫描结果摘要"""
    print_separator()
    print("快速扫描结果摘要:")
    print_separator()
    
    # 提取关键信息
    scan_id = report.get("scan_id", "N/A")
    target = report.get("target_url", "N/A")
    status = report.get("status", "N/A")
    
    # 处理不同格式的报告
    vulnerability_count = report.get("vulnerability_count", 0)
    
    # 检查是否有漏洞列表
    vulnerabilities = report.get("vulnerabilities", [])
    if isinstance(vulnerabilities, list) and len(vulnerabilities) > 0:
        vulnerability_count = len(vulnerabilities)
    
    # 显示基本信息
    print(f"扫描ID: {scan_id}")
    print(f"目标URL: {target}")
    print(f"扫描状态: {status}")
    
    # 显示时间信息（如果有）
    if "created_at" in report:
        print(f"创建时间: {report['created_at']}")
    if "completed_at" in report:
        print(f"完成时间: {report['completed_at']}")
    
    print(f"发现漏洞数: {vulnerability_count}")
    
    # 显示简单的风险评估
    risk_level = "低"
    if vulnerability_count >= 5:
        risk_level = "高"
    elif vulnerability_count >= 1:
        risk_level = "中"
    
    print(f"\n整体风险等级: {risk_level}")
    
    # 简单的建议
    if vulnerability_count > 0:
        print("\n安全建议:")
        print("1. 详细审查扫描报告中的所有漏洞")
        print("2. 优先修复高风险漏洞")
        print("3. 实施定期安全扫描策略")

def quick_demo():
    """快速演示主函数"""
    print_header("漏洞扫描器快速演示")
    print("此演示使用优化的扫描配置，提供更快的执行速度")
    print("适用于快速验证扫描器功能")
    print("注意：仅在授权的测试环境中运行此脚本")
    print_separator()
    
    # 检查API健康状态
    print("检查API服务状态...")
    if not check_api_health():
        print("❌ API服务不可用，请确保扫描器服务正在运行")
        print("请先启动扫描器服务：python start.py")
        return
    
    print("✅ API服务正常运行")
    print_separator()
    
    # 显示可用的快速目标列表
    print("可用的快速扫描目标:")
    for i, target in enumerate(QUICK_TARGETS):
        print(f"{i+1}. {target['name']} - {target['description']}")
    
    # 询问用户选择哪个目标
    try:
        choice = input("\n请选择要扫描的目标编号 (1-3, 默认为1): ")
        if not choice:
            choice = 1
        else:
            choice = int(choice)
            
        if 1 <= choice <= len(QUICK_TARGETS):
            selected_target = QUICK_TARGETS[choice - 1]
            perform_quick_scan(selected_target)
        else:
            print("无效的选择，将使用默认目标")
            perform_quick_scan(QUICK_TARGETS[0])
    except ValueError:
        print("无效的输入，将使用默认目标")
        perform_quick_scan(QUICK_TARGETS[0])
    
    print("\n快速演示完成")
    print("感谢使用漏洞扫描器!")

if __name__ == "__main__":
    quick_demo()